package com.rizzo.back.test.dao;


import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlWriter;
import org.dbunit.operation.DatabaseOperation;
import org.springframework.context.ApplicationContext;
import org.springframework.core.io.Resource;

import javax.sql.DataSource;
import java.io.*;
import java.sql.Connection;
import java.sql.SQLException;
import java.text.MessageFormat;
import java.util.List;

/**
 * Utility class for using DbUnit
 * Contains methods for clearing and filling the database,
 * creating schemas, a DTD or exporting data.
 */
public class DBUnitUtil {

	/**
	 * Fill the database with data from a resource file.
	 * @param dataFile the .xml resource containing the dataset
	 * @param dataSource the datasource to write to
	 * @throws SQLException in case of SQL errors
	 * @throws IOException in case of IO errors
	 */
	public static void fillDb(File dataFile, DataSource dataSource)
			throws SQLException, IOException {
		fillDb(new FileInputStream(dataFile), dataSource);
	}

	/**
	 * Fill the database with data from an inputStream
	 * @param inputStream the stream
	 * @param dataSource the datasource to write to
     * @throws SQLException in case of SQL errors
     * @throws IOException in case of IO errors
	 */
	public static void fillDb(InputStream inputStream, DataSource dataSource)
			throws SQLException, IOException {
		Connection connection = dataSource.getConnection();
		try {
			FlatXmlDataSet dataSet = new FlatXmlDataSet(inputStream);
			IDatabaseConnection dbUnitConn = new DatabaseConnection(connection);
            dbUnitConn.getConfig().setFeature(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, true);
			DatabaseOperation.CLEAN_INSERT.execute(dbUnitConn, dataSet);
		} catch (DataSetException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} catch (DatabaseUnitException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} finally {
			connection.close();
		}
	}

	/**
	 * Removes all data from the database
	 * @param dataFile the data file used for defining tables.
	 * @param dataSource the datasource
     * @throws SQLException in case of SQL errors
     * @throws IOException in case of IO errors
	 */
	public static void clearDb(File dataFile, DataSource dataSource)
			throws SQLException, IOException {
		clearDb(new FileInputStream(dataFile), dataSource);
	}

	/**
	 * Removes all data from the database
	 * @param inputStream the stream
	 * @param dataSource the datasource to write to
     * @throws SQLException in case of SQL errors
     * @throws IOException in case of IO errors
	 */
	public static void clearDb(InputStream inputStream, DataSource dataSource)
			throws SQLException, IOException {
		Connection connection = dataSource.getConnection();
		try {
			FlatXmlDataSet dataSet = new FlatXmlDataSet(inputStream);
			IDatabaseConnection dbUnitConn = new DatabaseConnection(connection);
            dbUnitConn.getConfig().setFeature(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, true);
			dbUnitConn.getConnection().createStatement().executeUpdate(
					"SET REFERENTIAL_INTEGRITY FALSE");
			DatabaseOperation.DELETE_ALL.execute(dbUnitConn, dataSet);
			dbUnitConn.getConnection().createStatement().executeUpdate(
					"SET REFERENTIAL_INTEGRITY TRUE");
		} catch (DataSetException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} catch (DatabaseUnitException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} finally {
			connection.close();
		}
	}

	/**
	 * Creates a DTD for a given data resource
	 * @param dataSource the data source
	 * @return the dtd
     * @throws SQLException in case of SQL errors
	 */
	public static String createDtd(DataSource dataSource) throws SQLException {
		Connection connection = dataSource.getConnection();
		try {
			IDatabaseConnection dbUnitConn = new DatabaseConnection(connection);
            dbUnitConn.getConfig().setFeature(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, true);
			// write DTD file
			StringWriter stringWriter = new StringWriter();
	        FlatDtdDataSet.write(dbUnitConn.createDataSet(), stringWriter);
	        return stringWriter.toString();
		} catch (DataSetException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} catch (IOException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} catch (DatabaseUnitException e) {
            SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
        } finally {
			connection.close();
		}
	}

	/**
	 * Exports the data in the DB to a given file
	 * @param dataSource the datasource to read
	 * @param file the file to write.
     * @throws SQLException in case of SQL errors
     * @throws IOException in case of IO errors
	 */
	public static void exportDb(DataSource dataSource, File file)
		throws SQLException, IOException {
		exportDb(dataSource, new FileOutputStream(file));
	}

	/**
	 * Exports the data in the DB to a given file
	 * @param dataSource the datasource to read
	 * @param stream the stream to write to.
     * @throws SQLException in case of SQL errors
     * @throws IOException in case of IO errors
	 */
	public static void exportDb(DataSource dataSource, OutputStream stream) throws SQLException, IOException {
		Connection connection = dataSource.getConnection();
		try {
			IDatabaseConnection dbUnitConn = new DatabaseConnection(connection);
            dbUnitConn.getConfig().setFeature(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, true);
			// write DTD file
			FlatXmlWriter datasetWriter = new FlatXmlWriter(stream);
			datasetWriter.write(dbUnitConn.createDataSet());
		} catch (DataSetException e) {
			SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
		} catch (DatabaseUnitException e) {
            SQLException exception = new SQLException();
			exception.initCause(e);
			throw exception;
        } finally {
			connection.close();
		}
	}

	/**
	 * Get the location of the DBUnit data file with test data.
	 *
	 * The default location is a classpath location and is determined by using
	 * the name of a class. When the <code>testClass</code> is given at
	 * construction time, this class will be used to determine the name of the
	 * XML file, otherwise the class name of this test case will be used.
	 *
     * @param clazz the test clazz to get a data file for
     * @param context the current application context
	 * @return the location of the DBUnit data file with test data.
	 */
	public static File getDataFile(Class<?> clazz, ApplicationContext context) {
		String fileName = clazz.getName().replaceAll("\\.", "/")
				+ ".xml";
        Resource resource = context.getResource("classpath:" + fileName);
		if (!resource.exists()) {
            resource = context.getResource("classpath:com/rizzo/back/test/dao/DefaultTestData.xml");
		}
        File file = null;
        try {
            file = resource.getFile();
        } catch (IOException e) {
            //TODO empty catch
        }
        return file;
    }

    /**
     * Creates schemas in a given a datasource
     * @param dataSource the data source to create schemas in
     * @param createStatement The statement to create the schema ({0} is placeholder for by the schema name)
     * @param schemaNames the list of schema names
     * @return the data source
     * @throws SQLException if creating the schemas failed
     */
    public static DataSource dataSourceWithSchemas(DataSource dataSource, String createStatement, List<String> schemaNames) throws SQLException {
        Connection connection = dataSource.getConnection();
        for (String schemaName : schemaNames) {
            connection.createStatement().executeUpdate(
                    new MessageFormat(createStatement).format(new Object [] {schemaName})
                );
        }
        connection.close();
        return dataSource;
    }
}